#! -*- coding: utf-8 -*-

import os
import json
import numpy as np
from tqdm import tqdm
from typing import Dict, List


import tensorflow as tf
from tensorflow.keras.layers import (Input, Lambda, Embedding, LSTM,
        Dense, Conv2D, MaxPooling2D, AveragePooling2D, ZeroPadding2D,
        Flatten, Activation, GlobalAveragePooling2D, GlobalMaxPooling2D,
        add, Layer, InputSpec, BatchNormalization, Concatenate, LeakyReLU)

from tensorflow.keras import backend as K
from tensorflow.keras.models import Model
from tensorflow.keras import initializers
from tensorflow.keras.callbacks import Callback
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import get_source_inputs


from autoTitleSeq2seq.model_layers import ScaleShift, OurBidirectional, \
    SelfModulatedLayerNormalization, Attention, LayerNormalization

from autoTitleSeq2seq.layer_process import LayerProcess
from utils.io import write_json

# ---文件路径
PACKAGE_PATH = os.path.dirname(os.path.dirname(__file__))
print(PACKAGE_PATH)
data_path = os.path.join(PACKAGE_PATH, 'data/auto_title/lcsts_data_short.json')  # 数据集：短文本
statis_voc = os.path.join(PACKAGE_PATH, 'autoTitleSeq2seq/common/seq2seq_voc.json')   # 词表（包括自定义字符）

s1 = u'夏天来临，皮肤在强烈紫外线的照射下，晒伤不可避免，因此，晒后及时修复显得尤为重要，' \
     u'否则可能会造成长期伤害。专家表示，选择晒后护肤品要慎重，芦荟凝胶是最安全，有效的一种选择，' \
     u'晒伤严重者，还请及时就医 。'
s2 = u'8月28日，网络爆料称，华住集团旗下连锁酒店用户数据疑似发生泄露。从卖家发布的内容看，数据包含华住旗下汉庭、' \
     u'禧玥、桔子、宜必思等10余个品牌酒店的住客信息。泄露的信息包括华住官网注册资料、酒店入住登记的身份信息及酒店开房记录，' \
     u'住客姓名、手机号、邮箱、身份证号、登录账号密码等。卖家对这个约5亿条数据打包出售。' \
     u'第三方安全平台威胁猎人对信息出售者提供的三万条数据进行验证，认为数据真实性非常高。当天下午，华住集 团发声明称，' \
     u'已在内部迅速开展核查，并第一时间报警。当晚，上海警方消息称，接到华住集团报案，警方已经介入调查。'

# ---参数配置
MIN_COUNT = 32  # 词频大小
MAXLEN = 400  # 最大长度，包括<start>和<end>标记
BATCH_SIZE = 2  # 批大小 64
EPOCHS = 100  # 训练轮数
CHARZ_SIZE = 128  # token维度大小
Z_DIM = 128       #
STEP_PER_EPOCH = 1000  # 每轮的步数


class PreProcess:
    """数据预处理
    （1）数据集加载
    （2）自定义词表
    """

    def __init__(self):
        self.data: Dict = self.load_data()  # 加载数据集
        self.custom_token = {0: 'MASK', 1: 'UNK', 2: 'START', 3: 'END'}         # 查找表自定义占位符
        self.custom_token_id = {'MASK': 0, 'UNK': 1, 'START': 2, 'END': 3}      # 查找表自定义占位符
        self.chars = {}
        self.id2char = {}
        self.char2id = {}
        self.gpu_flag = False  # self.gpu_info()

    def gpu_info(self):
        """gpu信息"""
        import tensorflow as tf
        devices = tf.config.experimental_list_devices()
        _LOCAL_DEVICES = [x.name for x in devices]
        for one in _LOCAL_DEVICES:
            if 'GPU' or 'gpu' in one:
                return True
        return False

    @staticmethod
    def load_data():
        """加载数据"""
        with open(data_path, 'r', encoding='utf-8') as f:
            data_dict = json.load(f)
            print(f'数据集样本数: {len(data_dict)}')
            return data_dict

    def txt_id(self, txt, start_end=False) -> List[int]:
        """文本转整数id
        （1）未登陆词填充；
        （2）补上<start>和<end>标记（可选）
        """
        if start_end:
            ids = [self.char2id.get(c, self.custom_token_id['UNK']) for c in txt[:MAXLEN - 2]]
            ids = [self.custom_token_id['START']] + ids + [self.custom_token_id['END']]
        else:
            # 转化，不加<start>和<end>标记
            ids = [self.char2id.get(c, self.custom_token_id['UNK']) for c in txt[:MAXLEN]]
        return ids

    def id_txt(self, ids):
        """id转文字，找不到的用空字符代替"""
        return ''.join([self.id2char.get(i, '') for i in ids])

    def padding(self, x: List[List[int]]) -> List[List[int]]:
        """padding为固定长度
        （1）padding为batch内的最大长度；
        （2）padding的部分用的是mask掩码，而不是padding为未登录词；
        （3）所以，对于title，作为label，它的结尾是'MASK': 0，'END': 3占位符在mask占位符前面。
            对于样本，由于没有<START>和<END>标记，所以不存在上述现象
        """
        max_len = max([len(i) for i in x])
        new_line = [line + [self.custom_token_id['MASK']] * (max_len - len(line)) for line in x]
        return new_line

    def statis(self):
        """统计数据集中的所有文本，以字为单位，自定义词表
        占位符：
            0: mask     # 掩码
            1: unk      # 未登录词
            2: start    # 开始
            3: end      # 结束
        """
        # 如果已经存在，直接加载
        if os.path.exists(statis_voc):
            statis_dic = json.load(open(statis_voc))
            self.chars, self.id2char, self.char2id = statis_dic['chars'], statis_dic['id2char'], statis_dic['char2id']
            self.id2char = {int(i): j for i, j in self.id2char.items()}
            return self.chars, self.char2id, self.id2char
        else:
            chars = {}  # token：字统计
            for line in tqdm(self.data):
                for w in line['content']:
                    chars[w] = chars.get(w, 0) + 1
                for w in line['title']:
                    chars[w] = chars.get(w, 0) + 1
            self.chars = {i: j for i, j in chars.items() if j >= MIN_COUNT}
            self.id2char = {i + 4: j for i, j in enumerate(self.chars)}
            self.id2char.update(self.custom_token)
            self.char2id = {j: i for i, j in self.id2char.items()}
            write_json(statis_voc, {'chars': self.chars, 'id2char': self.id2char, 'char2id': self.char2id})

    def get_one_sample(self) -> List[np.array]:
        """获取单个样本，用做调试
        （1）读取行数据；
        （2）文本格式化，根据词表转换为id；
        （3）content作为样本，不加<start>和<end>标记，title作为标签，需加上<start>和<end>标记；
        """
        batch_size = 2

        X, Y = [], []
        for line in self.data:
            X.append(self.txt_id(line['content']))  # content作为样本，不加<start>和<end>标记
            Y.append(self.txt_id(line['title'], start_end=True))  # title作为标签，需加上<start>和<end>标记
            if len(X) == batch_size:
                X = np.array(self.padding(X))
                Y = np.array(self.padding(Y))
                yield [X, Y]
                X, Y = [], []

    def data_generator(self):
        """数据生成器"""

        # X为文章，Y为标题
        X, Y = [], []

        while True:
            for line in tqdm(self.data):
                X.append(self.txt_id(line['content']))  # 样本：不加入开始、结束占位符
                Y.append(self.txt_id(line['title'], start_end=True))  # 标签：加入开始、结束占位符
                if len(X) == BATCH_SIZE:
                    X = np.array(self.padding(X))
                    Y = np.array(self.padding(Y))
                    yield [X, Y], None
                    X, Y = [], []


class Seq2SeqAutoSummary:
    """seq2seq模型"""

    def __init__(self, pre_obj: PreProcess):
        self.pre_process = pre_obj  # 预处理对象
        self.layer_process = LayerProcess(self.pre_process)
        self.model = self.model_build()  # 构建模型

    def seq_avgpool(self, x):
        """seq是[None, seq_len, s_size]的格式，
        mask是[None, seq_len, 1]的格式，先除去mask部分，
        然后再做avgpooling。
        """
        seq, mask = x
        return K.sum(seq * mask, 1) / (K.sum(mask, 1) + 1e-6)

    def gen_sent(self, s, topk=3, maxlen=64):
        """beam search解码
        每次只保留topk个最优候选结果；
        如果topk=1，那么就是贪心搜索
        """
        xid = np.array([self.pre_process.txt_id(s)] * topk)  # 输入转id
        yid = np.array([[2]] * topk)  # 解码均以<start>开头，这里<start>的id为2
        scores = [0] * topk  # 候选答案分数
        for i in range(maxlen):  # 强制要求输出不超过maxlen字
            proba = self.model.predict([xid, yid])[:, i, 3:]  # 直接忽略<padding>、<unk>、<start>
            log_proba = np.log(proba + 1e-6)  # 取对数，方便计算
            arg_topk = log_proba.argsort(axis=1)[:, -topk:]  # 每一项选出topk
            _yid = []  # 暂存的候选目标序列
            _scores = []  # 暂存的候选目标序列得分
            if i == 0:
                for j in range(topk):
                    _yid.append(list(yid[j]) + [arg_topk[0][j] + 3])
                    _scores.append(scores[j] + log_proba[0][arg_topk[0][j]])
            else:
                for j in range(topk):
                    for k in range(topk):  # 遍历topk*topk的组合
                        _yid.append(list(yid[j]) + [arg_topk[j][k] + 3])
                        _scores.append(scores[j] + log_proba[j][arg_topk[j][k]])
                _arg_topk = np.argsort(_scores)[-topk:]  # 从中选出新的topk
                _yid = [_yid[k] for k in _arg_topk]
                _scores = [_scores[k] for k in _arg_topk]
            yid = np.array(_yid)
            scores = np.array(_scores)
            best_one = np.argmax(scores)
            if yid[best_one][-1] == 3:
                return self.pre_process.id_txt(yid[best_one])
        # 如果maxlen字都找不到<end>，直接返回
        return self.pre_process.id_txt(yid[np.argmax(scores)])

    def model_build(self):
        """构建seq2seq模型"""
        x_in = Input(shape=(None,))     # 输入：(btz, seq_len_x)
        y_in = Input(shape=(None,))     # 输入：(btz, seq_len_y)
        x, y = x_in, y_in
        z_dim = Z_DIM

        # ---掩码层：比较元素值是否大于0（掩码），得到的结果shape=(btz, seq_len, 1)----
        # 初始值shape: (batch_size, seq_len)
        # 将其拆解后可知进行了如下操作:
        # inputs_mask = k.expand_dims(inputs, 2)
        # 1: 增加维度 在shape=2处增加一个维度将原来(10, 40)->(10, 40, 1)
        # inputs_mask = k.greater(inputs_mask, 0)
        # 2: 将为0的值置为false，大于0的则为True
        # inputs_mask = k.cast(inputs_mask, np.float32)
        # 3: 转换类型，True转换为1，False转换为0
        x_mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x, 2), 0), 'float32'))(x)
        y_mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x, 2), 0), 'float32'))(y)

        # one-hot层：(batch_size, 1, len(char)+4)， 维度len(char)+4指的是字典长度+4是因为占位符
        # x_one_hot = Lambda(self.to_one_hot, name='x_ont_hot')([x, x_mask])
        x_one_hot = Lambda(self.layer_process.to_one_hot, name='x_ont_hot')([x, x_mask])

        # 缩放平移层：仅仅是对x_one_hot的数据进行了缩放平移处理,
        # shape保持不变, 因此在embedding时input_dim依然是len(chars)+4
        # scale: (batch_size, 1, len(char)+4) 对内部数据做了处理
        # 这里起到的实际作用是什么？
        x_prior = ScaleShift()(x_one_hot)

        # 嵌入层：就是特征转向量
        # input_dim 则限定了输入特征的数目在此处要求input_dim = len(char)+4 + 1，
        # output_dim决定了输出向量的维度
        # Embedding会新建一个 len(char)+4 + 1 x CHARZ_SIZE 大小的查找表，然后进行参数训练。
        # embedding结果: (batch_size, max_len, CHARZ_SIZE) 对batch_size组每个字进行了向量映射。
        embedding = Embedding(input_dim=len(self.pre_process.chars) + 4, output_dim=CHARZ_SIZE)
        x = embedding(x)

        # 针对于层的归一化操作, 常用于RNN层, 此处的CuDNNLSTM层继承了部分RNN层的类, 而且LSTM本身就是RNN层的变种之一；
        # 其他相似的如BatchNor, GroupNorm, InstanceNorm, SwitchableNorm等基本上是针对不同情境下的归一化操作；
        x = LayerNormalization()(x)

        # encoder层：双层双向LSTM，在tensorflow 2中CuDNNLSTM并入到了lstm中，
        # 如果使用CuDNNLSTM需要考虑版本问题, 导入tf.compat.v1.keras.layers.CuDNNLSTM包
        x = OurBidirectional(LSTM(z_dim // 2, return_sequences=True))([x, x_mask])
        x = LayerNormalization()(x)
        x = OurBidirectional(LSTM(z_dim // 2, return_sequences=True))([x, x_mask])

        # seq_max_pool shape: (batch_size, max_len)，此处通过x_mask将x里面的mask值去掉后，进行了max处理
        x_mask_seq = Lambda(self.layer_process.seq_maxpool)([x, x_mask])

        # decoder：双层单向LSTM
        # 从embedding到此处中间设计层处理后的shape均为(batch_size, max_len, char_size)只是对内部数据进行了处理。
        # 通常，在需要将各个隐层的结果作为下一层的输入时，选择设置 return_sequences=True
        y = embedding(y)
        print(f"y_embedding: {y}, shape: {y.shape}")
        y = SelfModulatedLayerNormalization(z_dim // 4)([y, x_mask_seq])
        y = LSTM(z_dim, return_sequences=True)(y)
        y = SelfModulatedLayerNormalization(z_dim // 4)([y, x_mask_seq])
        y = LSTM(z_dim, return_sequences=True)(y)
        y = SelfModulatedLayerNormalization(z_dim // 4)([y, x_mask_seq])
        print('kskskd:', y)

        # Attention交互
        xy = Attention(8, 16)([y, x, x, x_mask])
        xy = Concatenate()([y, xy])

        # 输出分类
        xy = Dense(CHARZ_SIZE)(xy)
        # Relu激活函数的目的是将目标张量内部小于0的值置为0, leaky Relu函数则是将于小于0的值除以固定参数此处是0.2
        xy = LeakyReLU(0.2)(xy)
        xy = Dense(len(self.pre_process.chars) + 4)(xy)
        xy = Lambda(lambda x: (x[0] + x[1]) / 2)([xy, x_prior])  # 与先验结果平均
        xy = Activation('softmax')(xy)

        # 交叉熵作为loss，但mask掉padding部分
        # cross_entropy = K.sparse_categorical_crossentropy(y_in[:, 1:], xy[:, :-1])
        # cross_entropy = K.sum(cross_entropy * y_mask[:, 1:, 0]) / K.sum(y_mask[:, 1:, 0])
        cross_entropy_func = lambda x: K.sparse_categorical_crossentropy(x[0][:, 1:], x[1][:, :-1])
        cross_entropy = Lambda(cross_entropy_func)([y_in, xy])
        cross_entropy = tf.reduce_sum(cross_entropy * y_mask[:, 1:, 0]) / tf.reduce_sum(y_mask[:, 1:, 0])

        model = Model([x_in, y_in], xy)
        model.add_loss(cross_entropy)
        model.compile(optimizer=Adam(1e-3))

        return model


class Evaluate(Callback):
    """评估"""

    def __init__(self, seq2seq_obj: Seq2SeqAutoSummary):
        super(Evaluate, self).__init__()
        self.lowest = 1e10
        self.seq2seq = seq2seq_obj

    def on_epoch_end(self, epoch, logs=None):
        # 包括acc和loss的日志，并且可选地包括val_loss和val_acc
        # 训练过程中观察一两个例子，显示标题质量提高的过程
        print(self.seq2seq.gen_sent(s1))
        print(self.seq2seq.gen_sent(s2))
        # 保存最优结果
        if logs['loss'] <= self.lowest:
            self.lowest = logs['loss']
            self.seq2seq.model.save_weights('./best_model.weights')


if __name__ == '__main__':
    # 预处理
    pre_obj = PreProcess()
    pre_obj.statis()

    # 测试
    sample = None
    example_path = 'example.npz'
    if not os.path.exists(example_path):
        data_sample: List[np.array] = pre_obj.get_one_sample()
        for sam in data_sample:
            print(sam)
            np.savez(example_path, x=sam[0], y=sam[1])
            sample = sam
            break
    else:
        sample = np.load(example_path, allow_pickle=True)

    # 模型搭建
    s2s_obj = Seq2SeqAutoSummary(pre_obj)
    evaluator = Evaluate(s2s_obj)

    # 模型训练
    s2s_obj.model.fit_generator(
        pre_obj.data_generator(),
        steps_per_epoch=STEP_PER_EPOCH,
        epochs=EPOCHS,
        callbacks=[evaluator])
